import io
import json
import math

import requests
import random

from PIL import Image

import global_config

# 域名地址
base_url = global_config.base_url

with open('./source/comfyui_apis/workflow_img2img_api.json', 'r', encoding='utf-8') as file:
    dir_api_body = json.load(file)

with open('./source/comfyui_apis/workflow_svd_fast_api.json', 'r', encoding='utf-8') as file:
    svd_api_body = json.load(file)

with open('./source/comfyui_apis/workflow_text2img_api.json', 'r', encoding='utf-8') as file:
    text2img_api_body = json.load(file)


# 发送生图任务，同步返回任务id
# @params:
# imagename input目录下名称 or 绝对路径
def text2img_workflow(prompts, negtive, width,height):
    url = f'{base_url}prompt'
    # dir_api_body['client_id'] = "client_id"
    text2img_api_body['prompt']['7']['inputs']['seed'] = generate_seed()
    # text2img_api_body['prompt']['7']['inputs']['steps'] = 15 # todo test
    text2img_api_body['prompt']['31']['inputs']['text'] = prompts
    text2img_api_body['prompt']['4']['inputs']['text'] = negtive
    text2img_api_body['prompt']['8']['inputs']['batch_size'] = 1
    text2img_api_body['prompt']['8']['inputs']['width'] = width
    text2img_api_body['prompt']['8']['inputs']['height'] = height
    print(prompts, text2img_api_body['prompt']['7']['inputs']['seed'])
    response = requests.post(url, json=text2img_api_body)
    data = response.json()
    prompt_id = data['prompt_id']
    print(response.text)
    return prompt_id


# 发送生图任务，同步返回任务id
# @params:
# imagename input目录下名称 or 绝对路径
def img2img_workflow(imagename, prompts, negtive, width,height):
    url = f'{base_url}prompt'
    # dir_api_body['client_id'] = "client_id"
    dir_api_body['prompt']['7']['inputs']['seed'] = generate_seed()
    dir_api_body['prompt']['7']['inputs']['steps'] = 15  # todo test
    dir_api_body['prompt']['29']['inputs']['image'] = imagename
    dir_api_body['prompt']['31']['inputs']['text'] = prompts
    dir_api_body['prompt']['4']['inputs']['text'] = negtive
    dir_api_body['prompt']['36']['inputs']['amount'] = 1
    print(prompts, dir_api_body['prompt']['7']['inputs']['seed'])
    response = requests.post(url, json=dir_api_body)
    data = response.json()
    prompt_id = data['prompt_id']
    print(response.text)
    return prompt_id


def svd_workflow(imagepath, index):
    url = f'{base_url}prompt'
    # dir_api_body['client_id'] = "client_id"
    svd_api_body['prompt']['3']['inputs']['seed'] = generate_seed()
    # svd_api_body['prompt']['3']['inputs']['steps'] = 3 #todo test
    # svd_api_body['prompt']['12']['inputs']['video_frames'] = 1 # todo test
    # svd_api_body['prompt']['57']['inputs']['directory'] = image_dir
    svd_api_body['prompt']['13']['inputs']['image'] = imagepath
    # svd_api_body['prompt']['57']['inputs']['start_index'] = index
    response = requests.post(url, json=svd_api_body)
    data = response.json()
    prompt_id = data['prompt_id']
    print(response.text)
    return prompt_id


# 获取详细任务队列信息，正在运行的以及挂起的
def get_queue():
    url = f'{base_url}queue'
    response = requests.get(url)
    data = response.json()
    # print(data['queue_pending'])
    return data['queue_running'], data['queue_pending']


# 获取生成的图片bytes
def get_view_image(filename):
    url = f'{base_url}view?filename={filename}'
    response = requests.get(url)
    image = Image.open(io.BytesIO(response.content))
    return image


# 获取生成的视频bytes
def get_view_video(filename):
    url = f'{base_url}view?filename={filename}'
    response = requests.get(url)
    if response.status_code == 200:
        return response.content
    return None


# 根据任务id，查询生成的图片信息
# return：image list
def get_history(prompt_id):
    url = f'{base_url}history/{prompt_id}'
    response = requests.get(url)
    data = response.json()
    if prompt_id in data:
        outputs = response.json()[prompt_id]['outputs']
        # 22 是图生图
        if '22' in outputs:
            images = outputs['22']['images']
            for imageinfo in images:
                if imageinfo['type'] != 'output':
                    return []
            return [v['filename'] for v in images]
        elif '35' in outputs:
            gifs = outputs['35']['gifs']
            for gif in gifs:
                if gif['type'] != 'output':
                    return ''
                return [gif['filename']]
    return []


#

# 生成随机种子
def generate_seed():
    return abs(random.randrange(start=-1125899906842624, stop=1125899906842624))
